import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.patches import Rectangle
import matplotlib.cm as cm
import os

# 设置中文字体
plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义网络层数据
layers = [
    {"name": "Conv2d", "in_channels": 2, "out_channels": 64, "in_shape": [512, 512], "out_shape": [256, 256]},
    {"name": "ResNet18-layer1", "in_channels": 64, "out_channels": 64, "in_shape": [256, 256], "out_shape": [256, 256]},
    {"name": "ResNet18-layer2", "in_channels": 64, "out_channels": 128, "in_shape": [256, 256], "out_shape": [128, 128]},
    {"name": "ResNet18-layer3", "in_channels": 128, "out_channels": 256, "in_shape": [128, 128], "out_shape": [64, 64]},
    {"name": "ResNet18-layer4", "in_channels": 256, "out_channels": 512, "in_shape": [64, 64], "out_shape": [32, 32]},
    {"name": "AdaptiveAvgPool2d", "in_channels": 512, "out_channels": 512, "in_shape": [8, 8], "out_shape": [1, 1]}
]

output_dir = "docs/feature_visualizations"
os.makedirs(output_dir, exist_ok=True)

# 为每个层创建单独的图像
for i, layer in enumerate(layers):
    # 创建图形
    fig = plt.figure(figsize=(12, 10))
    
    # 输入特征可视化
    ax1 = fig.add_subplot(121, projection='3d')
    
    # 只绘制部分通道以保持可视化清晰度
    channels_to_show = min(layer['in_channels'], 10)
    channel_step = max(1, layer['in_channels'] // channels_to_show)
    
    colors = cm.viridis(np.linspace(0, 1, channels_to_show))
    
    for c in range(channels_to_show):
        actual_channel = c * channel_step
        x = np.arange(layer['in_shape'][0])
        y = np.arange(layer['in_shape'][1])
        X, Y = np.meshgrid(x, y)
        
        # 为每个通道创建一个平面
        Z = np.ones_like(X) * actual_channel
        
        # 绘制通道平面
        ax1.plot_surface(X, Y, Z, rstride=10, cstride=10, color=colors[c], alpha=0.7)
    
    ax1.set_title(f"输入: {layer['in_channels']}通道 [{layer['in_shape'][0]}x{layer['in_shape'][1]}]")
    ax1.set_xlabel('宽度')
    ax1.set_ylabel('高度')
    ax1.set_zlabel('通道')
    ax1.set_xlim(0, layer['in_shape'][0])
    ax1.set_ylim(0, layer['in_shape'][1])
    ax1.set_zlim(0, layer['in_channels'])
    
    # 输出特征可视化
    ax2 = fig.add_subplot(122, projection='3d')
    
    # 只绘制部分通道以保持可视化清晰度
    channels_to_show = min(layer['out_channels'], 10)
    channel_step = max(1, layer['out_channels'] // channels_to_show)
    
    colors = cm.plasma(np.linspace(0, 1, channels_to_show))
    
    for c in range(channels_to_show):
        actual_channel = c * channel_step
        x = np.arange(layer['out_shape'][0])
        y = np.arange(layer['out_shape'][1])
        X, Y = np.meshgrid(x, y)
        
        # 为每个通道创建一个平面
        Z = np.ones_like(X) * actual_channel
        
        # 绘制通道平面
        ax2.plot_surface(X, Y, Z, rstride=2, cstride=2, color=colors[c], alpha=0.7)
    
    ax2.set_title(f"输出: {layer['out_channels']}通道 [{layer['out_shape'][0]}x{layer['out_shape'][1]}]")
    ax2.set_xlabel('宽度')
    ax2.set_ylabel('高度')
    ax2.set_zlabel('通道')
    ax2.set_xlim(0, layer['out_shape'][0])
    ax2.set_ylim(0, layer['out_shape'][1])
    ax2.set_zlim(0, layer['out_channels'])
    
    # 设置整个图形的标题
    fig.suptitle(f"{layer['name']} 特征维度变化", fontsize=16)
    
    # 调整布局
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])  # 为suptitle留出空间
    
    # 保存图像
    filename = os.path.join(output_dir, f"{layer['name']}_feature_dimensions.png")
    plt.savefig(filename, dpi=300, bbox_inches='tight')
    plt.close()  # 关闭图形以释放内存
    
    print(f"已保存 {layer['name']} 的可视化图像到 {filename}")

print(f"所有可视化图像已保存到 {output_dir} 目录")    